import numpy as np
from copy import deepcopy

import torch
import torch.nn as nn
import torch.nn.functional as F

# TAST
from typing import List
import copy

# DeYO
import torchvision
from einops import rearrange


@torch.jit.script
def softmax_entropy(x: torch.Tensor) -> torch.Tensor:
    """Entropy of softmax distribution from logits."""
    return -(x.softmax(1) * x.log_softmax(1)).sum(1)


def collect_params(model):
    """Collect the affine scale + shift parameters from batch norms.
    Walk the model's modules and collect all batch normalization parameters.
    Return the parameters and their names.
    Note: other choices of parameterization are possible!
    """
    params = []
    names = []
    for nm, m in model.named_modules():
        if isinstance(m, nn.BatchNorm2d):
            for np, p in m.named_parameters():
                if np in ['weight', 'bias']:  # weight is scale, bias is shift
                    params.append(p)
                    names.append(f"{nm}.{np}")
    return params, names


def copy_model_and_optimizer(model, optimizer):
    """Copy the model and optimizer states for resetting after adaptation."""
    model_state = deepcopy(model.state_dict())
    optimizer_state = deepcopy(optimizer.state_dict())
    return model_state, optimizer_state


def load_model_and_optimizer(model, optimizer, model_state, optimizer_state):
    """Restore the model and optimizer states from copies."""
    model.load_state_dict(model_state, strict=True)
    optimizer.load_state_dict(optimizer_state)


def configure_model(model):
    """Configure model for use with tent."""
    # train mode, because tent optimizes the model to minimize entropy
    model.train()
    # disable grad, to (re-)enable only what tent updates
    model.requires_grad_(False)
    # configure norm for tent updates: enable grad + force batch statisics
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.requires_grad_(True)
            # force use of batch stats in train and eval modes
            m.track_running_stats = False
            m.running_mean = None
            m.running_var = None
    return model


def check_model(model):
    """Check model for compatability with tent."""
    is_training = model.training
    assert is_training, "tent needs train mode: call model.train()"
    param_grads = [p.requires_grad for p in model.parameters()]
    has_any_params = any(param_grads)
    has_all_params = all(param_grads)
    assert has_any_params, "tent needs params to update: " \
                           "check which require grad"
    assert not has_all_params, "tent should not update all params: " \
                               "check which require grad"
    has_bn = any([isinstance(m, nn.BatchNorm2d) for m in model.modules()])
    assert has_bn, "tent needs normalization for its optimization"


class ERM(nn.Module):
    def __init__(self,model):
        super().__init__()
        self.model = model.eval()
        self.featurizer = model.featurizer
        self.classifier = model.classifier
        
    @torch.no_grad()
    def forward(self, x):
        outputs = self.model.predict(x)
        return outputs
        

class BN(nn.Module):
    def __init__(self, model, steps=1, episodic=False):
        super().__init__()
        self.model = model
        self.steps = steps
        self.episodic = episodic
        assert self.steps>=0, 'steps must be non-negative'
        if self.steps==0:
            self.model.eval()
    
    @torch.no_grad()
    def forward(self, x):
        if self.steps>0:
            for _ in range(self.steps):
                outputs = self.model.predict(x)
        else:
            outputs = self.model.predict(x)
        return outputs


class Tent(nn.Module):
    """
    ICLR,2021
    Tent adapts a model by entropy minimization during testing.
    Once tented, a model adapts itself by updating on every forward.
    """
    def __init__(self, model, optimizer, steps=1, episodic=False):
        super().__init__()
        self.model = model
        self.optimizer = optimizer
        self.steps = steps
        assert steps > 0, "tent requires >= 1 step(s) to forward and update"
        self.episodic = episodic

        # note: if the model is never reset, like for continual adaptation,
        # then skipping the state copy would save memory
        self.model_state, self.optimizer_state = \
            copy_model_and_optimizer(self.model, self.optimizer)

    def forward(self, x):
        if self.episodic:
            self.reset()

        for _ in range(self.steps):
            outputs = self.forward_and_adapt(x, self.model, self.optimizer)

        return outputs

    def reset(self):
        if self.model_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer state")
        load_model_and_optimizer(self.model, self.optimizer,
                                 self.model_state, self.optimizer_state)
    
    @torch.enable_grad()
    def forward_and_adapt(self,x, model, optimizer):
        """Forward and adapt model on batch of data.
        Measure entropy of the model prediction, take gradients, and update params.
        """
        # forward
        outputs = model.predict(x)
        # adapt
        loss = softmax_entropy(outputs).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()        
        return outputs


class PseudoLabel(nn.Module):
    def __init__(self, model, optimizer, beta=0.9,steps=1, episodic=False):
        super().__init__()
        self.model = model
        self.optimizer = optimizer
        self.steps = steps
        assert steps > 0, "requires >= 1 step(s) to forward and update"
        self.episodic = episodic
        self.beta = beta  #threshold for selecting pseudo labels
        self.model_state, self.optimizer_state = \
            copy_model_and_optimizer(self.model, self.optimizer)

    def forward(self, x):
        if self.episodic:
            self.reset()

        for _ in range(self.steps):
            outputs = self.forward_and_adapt(x, self.model, self.optimizer)

        return outputs

    def reset(self):
        if self.model_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer state")
        load_model_and_optimizer(self.model, self.optimizer,
                                 self.model_state, self.optimizer_state)
    
    @torch.enable_grad()
    def forward_and_adapt(self,x, model, optimizer):
        # forward
        outputs = model.predict(x)
        # adapt        
        scores = F.softmax(outputs,1)
        py,y_prime = torch.max(scores,1)
        mask = py > self.beta
        loss = F.cross_entropy(outputs[mask],y_prime[mask])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return outputs


class SHOTIM(nn.Module):
    """
    SHOT-IM ,ICML 2020
    """
    def __init__(self, model, optimizer, steps=1, episodic=False):
        super().__init__()
        self.model = model
        self.optimizer = optimizer
        self.steps = steps
        assert steps > 0, "requires >= 1 step(s) to forward and update"
        self.episodic = episodic

        # note: if the model is never reset, like for continual adaptation,
        # then skipping the state copy would save memory
        self.model_state, self.optimizer_state = \
            copy_model_and_optimizer(self.model, self.optimizer)

    def forward(self, x):
        if self.episodic:
            self.reset()

        for _ in range(self.steps):
            outputs = self.forward_and_adapt(x, self.model, self.optimizer)

        return outputs

    def reset(self):
        if self.model_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer state")
        load_model_and_optimizer(self.model, self.optimizer,
                                 self.model_state, self.optimizer_state)
    
    @torch.enable_grad()
    def forward_and_adapt(self,x, model, optimizer):
        # forward
        outputs = model.predict(x)
        # adapt
        ent_loss = softmax_entropy(outputs).mean(0)
        softmax_out = F.softmax(outputs, dim=1)
        msoftmax = softmax_out.mean(dim=0)
        div_loss = torch.sum(msoftmax * torch.log(msoftmax + 1e-5))
        loss = ent_loss + div_loss
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return outputs
    
class SHOT(nn.Module):
    """
    SHOT ,ICML 2020
    """
    def __init__(self, model, optimizer, steps=1, episodic=False):
        super().__init__()
        self.model = model
        self.optimizer = optimizer
        self.steps = steps
        assert steps > 0, "requires >= 1 step(s) to forward and update"
        self.episodic = episodic

        # note: if the model is never reset, like for continual adaptation,
        # then skipping the state copy would save memory
        self.model_state, self.optimizer_state = \
            copy_model_and_optimizer(self.model, self.optimizer)
        
        self.beta = 0.3

    def forward(self, x, y):
        if self.episodic:
            self.reset()

        for _ in range(self.steps):
            outputs = self.forward_and_adapt(x, self.model, self.optimizer)

        return outputs

    def reset(self):
        if self.model_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer state")
        load_model_and_optimizer(self.model, self.optimizer,
                                 self.model_state, self.optimizer_state)
    
    @torch.enable_grad()
    def forward_and_adapt(self,x, model, optimizer):
        # forward
        # outputs = model.predict(x)
        z = model.featurizer(x)
        outputs = model.classifier(z)
        print(z.shape)
        print(outputs.shape)
        # adapt
        ent_loss = softmax_entropy(outputs).mean(0)
        softmax_out = F.softmax(outputs, dim=1)
        msoftmax = softmax_out.mean(dim=0)
        div_loss = torch.sum(msoftmax * torch.log(msoftmax + 1e-5))
        loss = ent_loss + div_loss

        # implemented by me
        # c0 = outputs.T @ z / torch.unsqueeze(torch.sum(outputs, dim=0),dim=-1)
        # y_hat = torch.argmax((F.normalize(c0) @ F.normalize(z).T),dim=0)  # [128]
        # c1 = F.one_hot(y_hat).float().T @ z / torch.unsqueeze(torch.sum(F.one_hot(y_hat).float(), dim=0),dim=-1)
        # y_hat = torch.argmax((F.normalize(c1) @ F.normalize(z).T),dim=0)
        # loss += nn.CrossEntropyLoss()(outputs, y_hat) * self.beta
        
        # implemented by TAST
        py, y_prime = F.softmax(outputs, dim=-1).max(1)
        flag = py > 0.9
        clf_loss = F.cross_entropy(outputs[flag], y_prime[flag])
        loss += clf_loss * self.beta

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return outputs
        
        
class T3A(nn.Module):
    """
    T3A, NeurIPS 2021
    """
    def __init__(self,model,filter_K=100,steps=1,episodic=False):
        super().__init__()
        self.model = model.eval()
        self.featurizer = model.featurizer
        self.classifier = model.classifier
        self.steps = steps
        assert steps > 0, "requires >= 1 step(s) to forward and update"
        self.episodic = episodic
        self.filter_K = filter_K
        
        warmup_supports = self.classifier.fc.weight.data
        self.num_classes = warmup_supports.size()[0]
        self.warmup_supports = warmup_supports
        warmup_prob = self.classifier(self.warmup_supports)
        self.warmup_ent = softmax_entropy(warmup_prob)
        self.warmup_labels = F.one_hot(warmup_prob.argmax(1), num_classes=self.num_classes).float()
        self.supports = self.warmup_supports.data
        self.labels = self.warmup_labels.data
        self.ent = self.warmup_ent.data
             
    @torch.no_grad() 
    def forward(self,x):
        z = self.featurizer(x)
        p = self.classifier(z)
        yhat = torch.nn.functional.one_hot(p.argmax(1), num_classes=self.num_classes).float()
        ent = softmax_entropy(p)

        self.supports = self.supports.to(z.device)
        self.labels = self.labels.to(z.device)
        self.ent = self.ent.to(z.device)
        self.supports = torch.cat([self.supports,z])
        self.labels = torch.cat([self.labels,yhat])
        self.ent = torch.cat([self.ent,ent])
        
        supports, labels = self.select_supports()
        supports = torch.nn.functional.normalize(supports, dim=1)
        weights = (supports.T @ (labels))
        
        return z @ torch.nn.functional.normalize(weights, dim=0)

    def select_supports(self):
        ent_s = self.ent
        y_hat = self.labels.argmax(dim=1).long()
        filter_K = self.filter_K
        if filter_K == -1:
            indices = torch.LongTensor(list(range(len(ent_s))))

        indices = []
        indices1 = torch.LongTensor(list(range(len(ent_s)))).cuda()
        for i in range(self.num_classes):
            _, indices2 = torch.sort(ent_s[y_hat == i])
            indices.append(indices1[y_hat==i][indices2][:filter_K])
        indices = torch.cat(indices)

        self.supports = self.supports[indices]
        self.labels = self.labels[indices]
        self.ent = self.ent[indices]
        
        return self.supports, self.labels
      

class TSD(nn.Module):
    """
    Test-time Self-Distillation (TSD)
    CVPR 2023
    """
    def __init__(self,model,optimizer,lam=0.1,filter_K=100,steps=1,episodic=False):
        super().__init__()
        self.model = model
        self.featurizer = model.featurizer
        self.classifier = model.classifier
        self.optimizer = optimizer
        self.steps = steps
        assert steps > 0, "requires >= 1 step(s) to forward and update"
        self.episodic = episodic
        self.filter_K = filter_K
        
        warmup_supports = self.classifier.fc.weight.data.detach()
        self.num_classes = warmup_supports.size()[0]
        self.warmup_supports = warmup_supports
        warmup_prob = self.classifier(self.warmup_supports)
        
        self.warmup_ent = softmax_entropy(warmup_prob)
        self.warmup_labels = F.one_hot(warmup_prob.argmax(1), num_classes=self.num_classes).float()
        self.warmup_scores = F.softmax(warmup_prob,1)
                
        self.supports = self.warmup_supports.data
        self.labels = self.warmup_labels.data
        self.ent = self.warmup_ent.data
        self.scores = self.warmup_scores.data
        self.lam = lam
        
        
    def forward(self,x):
        z = self.featurizer(x)
        p = self.classifier(z)
                       
        yhat = F.one_hot(p.argmax(1), num_classes=self.num_classes).float()
        ent = softmax_entropy(p)
        scores = F.softmax(p,1)

        with torch.no_grad():
            self.supports = self.supports.to(z.device)
            self.labels = self.labels.to(z.device)
            self.ent = self.ent.to(z.device)
            self.scores = self.scores.to(z.device)
            self.supports = torch.cat([self.supports,z])
            self.labels = torch.cat([self.labels,yhat])
            self.ent = torch.cat([self.ent,ent])
            self.scores = torch.cat([self.scores,scores])
        
            supports, labels = self.select_supports()
            supports = F.normalize(supports, dim=1)
            weights = (supports.T @ (labels))
                
        dist,loss = self.prototype_loss(z,weights.T,scores,use_hard=False)

        loss_local = topk_cluster(z.detach().clone(),supports,self.scores,p,k=3)
        loss += self.lam*loss_local
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return p

    def select_supports(self):
        ent_s = self.ent
        y_hat = self.labels.argmax(dim=1).long()
        filter_K = self.filter_K
        if filter_K == -1:
            indices = torch.LongTensor(list(range(len(ent_s))))

        indices = []
        indices1 = torch.LongTensor(list(range(len(ent_s)))).cuda()
        for i in range(self.num_classes):
            _, indices2 = torch.sort(ent_s[y_hat == i])
            indices.append(indices1[y_hat==i][indices2][:filter_K])
        indices = torch.cat(indices)

        self.supports = self.supports[indices]
        self.labels = self.labels[indices]
        self.ent = self.ent[indices]
        self.scores = self.scores[indices]
        
        return self.supports, self.labels
    
    def prototype_loss(self,z,p,labels=None,use_hard=False,tau=1):
        #z [batch_size,feature_dim]
        #p [num_class,feature_dim]
        #labels [batch_size,]        
        z = F.normalize(z,1)
        p = F.normalize(p,1)
        dist = z @ p.T / tau
        if labels is None:
            _,labels = dist.max(1)
        if use_hard:
            """use hard label for supervision """
            #_,labels = dist.max(1)  #for prototype-based pseudo-label
            labels = labels.argmax(1)  #for logits-based pseudo-label
            loss =  F.cross_entropy(dist,labels)
        else:
            """use soft label for supervision """
            loss = softmax_kl_loss(labels.detach(),dist).sum(1).mean(0)  #detach is **necessary**
            #loss = softmax_kl_loss(dist,labels.detach()).sum(1).mean(0) achieves comparable results
        return dist,loss
    

class PROGRAM(nn.Module):
    def __init__(self,model,optimizer,lam=0.5,beta=0.04):
        super().__init__()
        self.model = model
        self.featurizer = model.featurizer
        self.classifier = model.classifier
        self.optimizer = optimizer
        self.lam = lam
        self.beta = beta
        self.eps = 1e-4


    def forward(self,x):
        warmup_supports = self.classifier.fc.weight.data.detach()
        self.num_classes = warmup_supports.size()[0]
        self.warmup_supports = torch.nn.functional.normalize(warmup_supports, p=2, dim=1)
        self.warmup_supports_T = self.warmup_supports.T.cuda()

        z = self.featurizer(x)
        p = self.classifier(z)

        # self.warmup_supports.shape  torch.Size([7, 512])
        # z.shape  torch.Size([128, 512])
        S = torch.matmul(z, self.warmup_supports_T)
        Z = torch.cat((torch.eye(self.num_classes).cuda(),S),dim=0)

        column_sums = torch.sum(Z, dim=0)
        D = torch.diag(column_sums).cuda()
        D_inv = torch.inverse(D).cuda()
        W = torch.matmul(Z, torch.matmul(D_inv, Z.T))
        rows, cols = W.size()
        Y_star=(1-self.lam)*torch.matmul(torch.inverse((torch.eye(rows,cols).cuda()-self.lam*W)),Z)
        Y_tilde = torch.softmax(Y_star[self.num_classes:, :], dim=1)
        
        y_argmax = torch.argmax(Y_tilde,dim=1)
        p_argmax = torch.argmax(p,dim=1)
        mask_same = (y_argmax == p_argmax)
        mask_diff = (y_argmax != p_argmax)
        loss1 = nn.CrossEntropyLoss(reduction='sum')(p[mask_same],y_argmax[mask_same])
        loss2 = self.symmetric_cross_entropy(p[mask_diff],Y_star[self.num_classes:, :][mask_diff]).sum(0)

        loss = (loss1 + loss2 * self.beta) / z.size()[0]

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        z = self.featurizer(x)
        p = self.classifier(z)

        return p
    
    def symmetric_cross_entropy(self,x,x_ema,alpha=0.5):  # implemented by RMT
        return -(1-alpha) * (x_ema.softmax(1) * x.log_softmax(1)).sum(1) - alpha * (x.softmax(1) * x_ema.log_softmax(1)).sum(1)

class PLL(nn.Module):
    def __init__(self,model,optimizer, margin_select=0.4, margin_end=0.2, margin_descend=0.01, temp_select=3, temp_loss=5, buffer_size=64):
        super().__init__()
        self.model = model
        self.featurizer = model.featurizer
        self.classifier = model.classifier
        self.optimizer = optimizer
        
        self.num_classes = self.classifier.fc.weight.data.size()[0]

        self.margin_select = margin_select
        self.margin_end = margin_end
        self.margin_descend = margin_descend
        self.temp_select = temp_select
        self.temp_loss = temp_loss
        self.buffer_size = buffer_size

        self.samples_buffer = None
        self.labels_buffer = None
    
    def forward(self,samples,labels):

        origin_sample_num = samples.shape[0]

        if self.samples_buffer != None:
            samples = torch.cat((samples,self.samples_buffer),dim=0)
            labels = torch.cat((labels,self.labels_buffer),dim=0)

        logits = self.model.predict(samples)

        probs_t = F.softmax(logits/self.temp_select,1)
        probs_de_t, _ = torch.sort(probs_t, descending=True)

        margins = probs_de_t[:,0] - probs_de_t[:,1]

        mask_hard = margins > self.margin_select
        mask_unselect = (probs_de_t[:,0] - probs_de_t[:,-1]) < self.margin_select
        mask_partial = ~ (mask_hard | mask_unselect)

        mask_unselect = mask_unselect.cpu()

        _, idxs = torch.sort(margins[mask_unselect], descending=True)
        idxs = idxs.cpu()
        if idxs.shape[0] > self.buffer_size:
            self.samples_buffer = samples[mask_unselect][idxs][:self.buffer_size]
            self.labels_buffer = labels[mask_unselect][idxs][:self.buffer_size]
        else:
            self.samples_buffer = samples[mask_unselect]
            self.labels_buffer = labels[mask_unselect]
        
        partial_labels = ((probs_t[mask_partial] + self.margin_select) > probs_de_t[mask_partial][:,0].reshape(-1,1)).long()    

        loss_hard = nn.CrossEntropyLoss(label_smoothing=0.0)(logits[mask_hard] / self.temp_loss, logits[mask_hard].detach().argmax(1))
        loss_partial = cc_loss(logits[mask_partial], partial_labels.detach(), self.temp_loss)

        mask_hard = mask_hard.cpu()
        mask_partial = mask_partial.cpu()
        # print('--------------------------')
        # print('batch num:',len(mask_hard))
        # print('hard num:',sum(mask_hard.long()).item())
        # print('partial num:',sum(mask_partial.long()).item())
        # hard_acc = sum((logits.argmax(1).cpu()[mask_hard] == labels[mask_hard]).long()) / sum(mask_hard.long())
        # print('hard label acc:{:.3f}'.format(hard_acc.item()))
        # partial_acc = torch.sum(partial_labels.cpu() & F.one_hot(labels[mask_partial], num_classes=self.num_classes)) / sum(mask_partial.long())
        # print('partial acc:{:.3f}'.format(partial_acc.item()))
        # avg_pl = torch.sum(partial_labels) / sum(mask_partial.long())
        # print('avg partial label:{:.3f}'.format(avg_pl.item()))
        # print('num classes:',self.num_classes)
        # print('margin select:',round(self.margin_select,3))
        # print('buffer num:',self.samples_buffer.shape[0])

        if sum(mask_hard.long()) == 0 and sum(mask_partial.long()) == 0:
            #print('total loss:0')
            return logits[0:origin_sample_num]
        elif sum(mask_hard.long()) != 0 and sum(mask_partial.long()) == 0:
            loss = loss_hard
            #print('hard label loss:{:.3f}'.format(loss_hard))
            #print('total loss:{:.3f}'.format(loss))
        elif sum(mask_hard.long()) == 0 and sum(mask_partial.long()) != 0:
            loss = loss_partial
            #print('partial label loss:{:.3f}'.format(loss_partial))
            #print('total loss:{:.3f}'.format(loss))
        else:
            lam_hard = sum(mask_hard.long()) / (sum(mask_hard.long()) + sum(mask_partial.long()))
            loss = loss_hard * lam_hard + loss_partial * (1 - lam_hard)
            #print('hard label loss:{:.3f}'.format(loss_hard))
            #print('partial label loss:{:.3f}'.format(loss_partial))
            #print('total loss:{:.3f}'.format(loss))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        if self.margin_select > self.margin_end:
            self.margin_select -= self.margin_descend

        return logits[0:origin_sample_num]
    
class PLL_imbalance(nn.Module):
    def __init__(self,model,optimizer, margin_select=0.4, margin_end=0.2, margin_descend=0.01, temp_select=3, temp_loss=5, buffer_size=64):
        super().__init__()
        self.model = model
        self.featurizer = model.featurizer
        self.classifier = model.classifier
        self.optimizer = optimizer
        
        self.num_classes = self.classifier.fc.weight.data.size()[0]

        self.margin_select = margin_select
        self.margin_end = margin_end
        self.margin_descend = margin_descend
        self.temp_select = temp_select
        self.temp_loss = temp_loss
        self.buffer_size = buffer_size

        self.samples_buffer = None
        self.labels_buffer = None

        # prototype(T3A)

        self.model_copy = copy.deepcopy(self.model).eval()

        self.filter_K = 100
        
        warmup_supports = self.model_copy.classifier.fc.weight.data
        # self.num_classes = warmup_supports.size()[0]
        self.warmup_supports = warmup_supports
        warmup_prob = self.model_copy.classifier(self.warmup_supports)
        self.warmup_ent = softmax_entropy(warmup_prob)
        self.warmup_labels = F.one_hot(warmup_prob.argmax(1), num_classes=self.num_classes).float()
        self.supports = self.warmup_supports.data
        self.labels = self.warmup_labels.data
        self.ent = self.warmup_ent.data

        self.hard_num = 0
        self.ours_num = 0
    
    def forward(self,samples,labels):

        self.origin_sample_num = samples.shape[0]  # 这里T3A要排除多余样本的影响

        if self.samples_buffer != None:
            samples = torch.cat((samples,self.samples_buffer),dim=0)
            labels = torch.cat((labels,self.labels_buffer),dim=0)

        logits = self.model.predict(samples)

        probs_t = F.softmax(logits/self.temp_select,1)
        probs_de_t, _ = torch.sort(probs_t, descending=True)

        margins = probs_de_t[:,0] - probs_de_t[:,1]

        mask_hard = margins > self.margin_select
        mask_unselect = (probs_de_t[:,0] - probs_de_t[:,-1]) < self.margin_select
        mask_partial = ~ (mask_hard | mask_unselect)

        mask_unselect = mask_unselect.cpu()

        _, idxs = torch.sort(margins[mask_unselect], descending=True)
        idxs = idxs.cpu()
        if idxs.shape[0] > self.buffer_size:
            self.samples_buffer = samples[mask_unselect][idxs][:self.buffer_size]
            self.labels_buffer = labels[mask_unselect][idxs][:self.buffer_size]
        else:
            self.samples_buffer = samples[mask_unselect]
            self.labels_buffer = labels[mask_unselect]
        
        partial_labels = ((probs_t[mask_partial] + self.margin_select) > probs_de_t[mask_partial][:,0].reshape(-1,1)).long()

        # label filter
        
        label_prototype = self.get_label_with_prototype(samples).argmax(1)
        mask_hard_same = logits[mask_hard].argmax(1) == label_prototype[mask_hard]
        rows = torch.arange(mask_partial.long().sum())
        cols = label_prototype[mask_partial]
        mask_partial_same = (partial_labels)[rows,cols] > 0  # label_prototype in partial labels

        # version 1
        # logits_final = torch.cat((logits[mask_hard][mask_hard_same],logits[mask_partial][mask_partial_same]),dim=0)
        # labels_final = torch.cat((label_prototype[mask_hard][mask_hard_same],label_prototype[mask_partial][mask_partial_same]))
        # loss = nn.CrossEntropyLoss()(logits_final / self.temp_loss, labels_final)

        loss_hard = nn.CrossEntropyLoss()(logits[mask_hard][mask_hard_same] / self.temp_loss, logits[mask_hard][mask_hard_same].detach().argmax(1))
        loss_partial = cc_loss(logits[mask_partial][mask_partial_same], partial_labels[mask_partial_same].detach(), self.temp_loss)
        lam_hard = mask_hard_same.long().sum() / (mask_hard_same.long().sum() + mask_partial_same.long().sum())
        loss = loss_hard * lam_hard + loss_partial * (1 - lam_hard)
        #loss = loss_hard

        self.hard_num += mask_hard.long().sum().item()
        self.ours_num += (mask_hard.long().sum().item() + mask_partial.long().sum().item())
        print("hard",self.hard_num)
        print("ours",self.ours_num)

        # print('--------------------------')
        # print('origin hard:', mask_hard.long().sum().item())
        # print('select hard:', mask_hard_same.long().sum().item())
        # print('origin partial:', mask_partial.long().sum().item())
        # print('select partial:', mask_partial_same.long().sum().item())

        
        mask_hard = mask_hard.cpu()
        mask_partial = mask_partial.cpu()
        mask_hard_same = mask_hard_same.cpu()
        mask_partial_same = mask_partial_same.cpu()
        
        # print('batch num:',len(mask_hard))
        # print('hard num:',mask_hard_same.long().sum().item())
        # print('partial num:',mask_partial_same.long().sum().item())
        # hard_acc = (logits.argmax(1).cpu()[mask_hard][mask_hard_same] == labels[mask_hard][mask_hard_same]).long().sum() / mask_hard_same.long().sum()
        # print('hard label acc:{:.3f}'.format(hard_acc.item()))
        # partial_acc = (partial_labels[mask_partial_same].cpu() & F.one_hot(labels[mask_partial][mask_partial_same], num_classes=self.num_classes)).sum() / mask_partial_same.long().sum()
        # print('partial acc:{:.3f}'.format(partial_acc.item()))
        # avg_pl = partial_labels[mask_partial_same].sum() / mask_partial_same.long().sum()
        # print('avg partial label:{:.3f}'.format(avg_pl.item()))
        # print('num classes:',self.num_classes)
        # print('margin select:',round(self.margin_select,3))
        # print('buffer num:',self.samples_buffer.shape[0])

        # if sum(mask_hard.long()) == 0 and sum(mask_partial.long()) == 0:
        #     print('total loss:0')
        #     return logits[0:origin_sample_num]
        # elif sum(mask_hard.long()) != 0 and sum(mask_partial.long()) == 0:
        #     loss = loss_hard
        #     print('hard label loss:{:.3f}'.format(loss_hard))
        #     print('total loss:{:.3f}'.format(loss))
        # elif sum(mask_hard.long()) == 0 and sum(mask_partial.long()) != 0:
        #     loss = loss_partial
        #     print('partial label loss:{:.3f}'.format(loss_partial))
        #     print('total loss:{:.3f}'.format(loss))
        # else:
        #     lam_hard = sum(mask_hard.long()) / (sum(mask_hard.long()) + sum(mask_partial.long()))
        #     loss = loss_hard * lam_hard + loss_partial * (1 - lam_hard)
        #     print('hard label loss:{:.3f}'.format(loss_hard))
        #     print('partial label loss:{:.3f}'.format(loss_partial))
        #     print('total loss:{:.3f}'.format(loss))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        if self.margin_select > self.margin_end:
            self.margin_select -= self.margin_descend

        # result = self.get_label_with_prototype(samples)
        # return result[0:self.origin_sample_num]

        return logits[0:self.origin_sample_num]
    
    @torch.no_grad() 
    def get_label_with_prototype(self,x):
        z = self.model_copy.featurizer(x)
        p = self.model_copy.classifier(z)
        yhat = torch.nn.functional.one_hot(p.argmax(1), num_classes=self.num_classes).float()
        ent = softmax_entropy(p)

        self.supports = self.supports.to(z.device)
        self.labels = self.labels.to(z.device)
        self.ent = self.ent.to(z.device)
        self.supports = torch.cat([self.supports,z[0:self.origin_sample_num,:]])
        self.labels = torch.cat([self.labels,yhat[0:self.origin_sample_num]])
        self.ent = torch.cat([self.ent,ent[0:self.origin_sample_num]])
        
        supports, labels = self.select_supports()
        supports = torch.nn.functional.normalize(supports, dim=1)
        weights = (supports.T @ (labels))
        
        return z @ torch.nn.functional.normalize(weights, dim=0)

    def select_supports(self):
        ent_s = self.ent
        y_hat = self.labels.argmax(dim=1).long()
        filter_K = self.filter_K
        if filter_K == -1:
            indices = torch.LongTensor(list(range(len(ent_s))))

        indices = []
        indices1 = torch.LongTensor(list(range(len(ent_s)))).cuda()
        for i in range(self.num_classes):
            _, indices2 = torch.sort(ent_s[y_hat == i])
            indices.append(indices1[y_hat==i][indices2][:filter_K])
        indices = torch.cat(indices)

        self.supports = self.supports[indices]
        self.labels = self.labels[indices]
        self.ent = self.ent[indices]
        
        return self.supports, self.labels
        
    
    # EMA implemented by RMT
    def ema_update_model(self, model_to_update, model_to_merge, momentum, device, update_all=False):
        if momentum < 1.0:
            for param_to_update, param_to_merge in zip(model_to_update.parameters(), model_to_merge.parameters()):
                if param_to_update.requires_grad or update_all:
                    param_to_update.data = momentum * param_to_update.data + (1 - momentum) * param_to_merge.data.to(device)
        return model_to_update

def reverse_ce(y, target):
    return - torch.mean(torch.sum(torch.log(1.0000001 - F.softmax(y, dim=1)) * (1 - target), dim=1))

def cc_loss(outputs, partialY, temp):
    sm_outputs = F.softmax(outputs / temp, dim=1)
    final_outputs = sm_outputs * partialY
    average_loss = - torch.log(final_outputs.sum(dim=1)).mean()
    return average_loss

def proden_loss(outputs, partialY, temp):
    probs = F.softmax(outputs / temp, dim=1)
    revisedY = partialY * probs
    revisedY = revisedY / (revisedY).sum(dim=1).repeat(revisedY.size(1), 1).transpose(0, 1)
    revisedY = revisedY.clone().detach()

    probs = F.softmax(outputs / temp, dim=1)

    l = revisedY * torch.log(probs)
    loss = (-torch.sum(l)) / l.size(0)
    return loss
        

def topk_labels(feature,supports,scores,k=3):
    feature = F.normalize(feature,1)
    supports = F.normalize(supports,1)
    sim_matrix = feature @ supports.T  #B,M
    _,idx_near = torch.topk(sim_matrix,k,dim=1)  #batch x K
    scores_near = scores[idx_near]  #batch x K x num_class
    soft_labels = torch.mean(scores_near,1)  #batch x num_class
    soft_labels = torch.argmax(soft_labels,1)
    return soft_labels
    

def topk_cluster(feature,supports,scores,p,k=3):
    #p: outputs of model batch x num_class
    feature = F.normalize(feature,1)
    supports = F.normalize(supports,1)
    sim_matrix = feature @ supports.T  #B,M
    topk_sim_matrix,idx_near = torch.topk(sim_matrix,k,dim=1)  #batch x K
    scores_near = scores[idx_near].detach().clone()  #batch x K x num_class
    diff_scores = torch.sum((p.unsqueeze(1) - scores_near)**2,-1)
    
    loss = -1.0* topk_sim_matrix * diff_scores
    return loss.mean()
    
    
def knn_affinity(X,knn):
    #x [N,D]
    N = X.size(0)
    X = F.normalize(X,1)
    dist = torch.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, p=2)  # [N, N]
    n_neighbors = min(knn + 1, N)
    knn_index = dist.topk(n_neighbors, -1, largest=False).indices[:, 1:]  # [N, knn]
    W = torch.zeros(N, N, device=X.device)
    W.scatter_(dim=-1, index=knn_index, value=1.0)
    return W
    
       
def softmax_mse_loss(input_logits, target_logits):
    """Takes softmax on both sides and returns MSE loss

    Note:
    - Returns the sum over all examples. Divide by the batch size afterwards
      if you want the mean.
    - Sends gradients to inputs but not the targets.
    """
    assert input_logits.size() == target_logits.size()
    input_softmax = F.softmax(input_logits, dim=1)
    target_softmax = F.softmax(target_logits, dim=1)

    mse_loss = (input_softmax-target_softmax)**2
    return mse_loss

    
def softmax_kl_loss(input_logits, target_logits):
    """Takes softmax on both sides and returns KL divergence

    Note:
    - Returns the sum over all examples. Divide by the batch size afterwards
      if you want the mean.
    - Sends gradients to inputs but not the targets.
    """
    assert input_logits.size() == target_logits.size()
    input_log_softmax = F.log_softmax(input_logits, dim=1)
    target_softmax = F.softmax(target_logits, dim=1)

    kl_div = F.kl_div(input_log_softmax, target_softmax, reduction='none')
    return kl_div        
        

def get_distances(X, Y, dist_type="cosine"):
    """
    Args:
        X: (N, D) tensor
        Y: (M, D) tensor
    """
    if dist_type == "euclidean":
        distances = torch.cdist(X, Y)
    elif dist_type == "cosine":
        distances = 1 - torch.matmul(F.normalize(X, dim=1), F.normalize(Y, dim=1).T)
    else:
        raise NotImplementedError(f"{dist_type} distance not implemented.")

    return distances


@torch.no_grad()
def soft_k_nearest_neighbors(features, features_bank, probs_bank):
    pred_probs = []
    K = 4
    for feats in features.split(64):
        distances = get_distances(feats, features_bank,"cosine")
        _, idxs = distances.sort()
        idxs = idxs[:, : K]
        # (64, num_nbrs, num_classes), average over dim=1
        probs = probs_bank[idxs, :].mean(1)
        pred_probs.append(probs)
    pred_probs = torch.cat(pred_probs)
    _, pred_labels = pred_probs.max(dim=1)

    return pred_labels, pred_probs



# TAST

def softmax_entropy(x: torch.Tensor) -> torch.Tensor:
    """Entropy of softmax distribution from logits."""
    return -(x.softmax(1) * x.log_softmax(1)).sum(1)

# motivated from "https://github.com/cs-giung/giung2/blob/c8560fd1b5/giung2/layers/linear.py"
class BatchEnsemble(nn.Module):
    def __init__(self, indim, outdim, ensemble_size, init_mode):
        super().__init__()
        self.ensemble_size = ensemble_size

        self.in_features = indim
        self.out_features = outdim

        # register parameters
        self.register_parameter(
            "weight", nn.Parameter(
                torch.Tensor(self.out_features, self.in_features)
            )
        )
        self.register_parameter(
            "bias", nn.Parameter(
                torch.Tensor(self.out_features)
            )
        )

        self.register_parameter(
            "alpha_be", nn.Parameter(
                torch.Tensor(self.ensemble_size, self.in_features)
            )
        )
        self.register_parameter(
            "gamma_be", nn.Parameter(
                torch.Tensor(self.ensemble_size, self.out_features)
            )
        )

        use_ensemble_bias = True
        if use_ensemble_bias and self.bias is not None:
            delattr(self, "bias")
            self.register_parameter("bias", None)
            self.register_parameter(
                "ensemble_bias", nn.Parameter(
                    torch.Tensor(self.ensemble_size, self.out_features)
                )
            )
        else:
            self.register_parameter("ensemble_bias", None)

        # initialize parameters
        self.init_mode = init_mode
        self.reset()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, D1 = x.size()
        r_x = x.unsqueeze(0).expand(self.ensemble_size, B, D1) #
        r_x = r_x.view(self.ensemble_size, -1, D1)
        r_x = r_x * self.alpha_be.view(self.ensemble_size, 1, D1)
        r_x = r_x.view(-1, D1)

        w_r_x = nn.functional.linear(r_x, self.weight, self.bias)

        _, D2 = w_r_x.size()
        s_w_r_x = w_r_x.view(self.ensemble_size, -1, D2)
        s_w_r_x = s_w_r_x * self.gamma_be.view(self.ensemble_size, 1, D2)
        if self.ensemble_bias is not None:
            s_w_r_x = s_w_r_x + self.ensemble_bias.view(self.ensemble_size, 1, D2)
        s_w_r_x = s_w_r_x.view(-1, D2)

        return s_w_r_x

    def reset(self):
        init_details = [0,1]
        initialize_tensor(self.weight, self.init_mode, init_details)
        initialize_tensor(self.alpha_be, self.init_mode, init_details)
        initialize_tensor(self.gamma_be, self.init_mode, init_details)
        if self.ensemble_bias is not None:
            initialize_tensor(self.ensemble_bias, "zeros")
        if self.bias is not None:
            initialize_tensor(self.bias, "zeros")

def initialize_tensor(
        tensor: torch.Tensor,
        initializer: str,
        init_values: List[float] = [],
    ) -> None:

    if initializer == "zeros":
        nn.init.zeros_(tensor)

    elif initializer == "ones":
        nn.init.ones_(tensor)

    elif initializer == "uniform":
        nn.init.uniform_(tensor, init_values[0], init_values[1])

    elif initializer == "normal":
        nn.init.normal_(tensor, init_values[0], init_values[1])

    elif initializer == "random_sign":
        with torch.no_grad():
            tensor.data.copy_(
                2.0 * init_values[1] * torch.bernoulli(
                    torch.zeros_like(tensor) + init_values[0]
                ) - init_values[1]
            )
    elif initializer == 'xavier_normal':
        torch.nn.init.xavier_normal_(tensor)

    elif initializer == 'kaiming_normal':
        torch.nn.init.kaiming_normal_(tensor)
    else:
        raise NotImplementedError(
            f"Unknown initializer: {initializer}"
        )

class TAST(nn.Module):
    def __init__(self, algorithm, lr=1e-3):
        super().__init__()
        # trained feature extractor and last linear classifier
        self.featurizer = algorithm.featurizer
        self.classifier = algorithm.classifier

        # store supports and corresponding labels
        warmup_supports = self.classifier.fc.weight.data.detach()
        self.warmup_supports = warmup_supports
        warmup_prob = self.classifier(self.warmup_supports)

        self.num_classes = warmup_supports.size()[0]
        self.warmup_labels = F.one_hot(warmup_prob.argmax(1), num_classes=self.num_classes).float()
        self.warmup_ent = softmax_entropy(warmup_prob)
        self.ent = self.warmup_ent.data

        self.supports = self.warmup_supports.data
        self.labels = self.warmup_labels.data

        # hparams
        # adapt_hparams_dict = {
        #     'num_ensemble': [1, 5, 10, 20],
        #     'filter_K': [1, 5, 20, 50, 100, -1],
        #     'gamma': [1, 3],
        #     'lr': [1e-3],
        #     'tau': [10],
        #     'k': [1, 2, 4, 8],
        #     'init_mode': ['kaiming_normal']
        # }
        self.filter_K = 50
        self.steps = 1
        self.num_ensemble = 10
        self.lr = lr
        self.tau = 10
        self.init_mode = 'kaiming_normal'
        self.k = 4

        self.cached_loader = False
        self.n_outputs = self.warmup_supports.shape[1]

        # multiple projection heads and its optimizer
        self.mlps = BatchEnsemble(self.n_outputs, self.n_outputs // 4, self.num_ensemble,
                                  self.init_mode).cuda()
        self.optimizer = torch.optim.Adam(self.mlps.parameters(), lr=self.lr)

    def forward(self, x, adapt=True):
        if not self.cached_loader:
            with torch.no_grad():
                z = self.featurizer(x)
        else:
            z = x

        if adapt:
            p_supports = self.classifier(z)
            yhat = torch.nn.functional.one_hot(p_supports.argmax(1), num_classes=self.num_classes).float()
            ent = softmax_entropy(p_supports)

            # prediction
            self.supports = self.supports.to(z.device)
            self.labels = self.labels.to(z.device)
            self.ent = self.ent.to(z.device)
            self.supports = torch.cat([self.supports, z])
            self.labels = torch.cat([self.labels, yhat])
            self.ent = torch.cat([self.ent, ent])

        supports, labels = self.select_supports()

        for _ in range(self.steps):
            p = self.forward_and_adapt(z, supports, labels)

        return p

    def compute_logits(self, z, supports, labels, mlp):
        '''
        :param z: unlabeled test examples
        :param supports: support examples
        :param labels: labels of support examples
        :param mlp: multiple projection heads
        :return: classification logits of z
        '''
        B, dim = z.size()
        N, dim_ = supports.size()

        mlp_z = mlp(z)
        mlp_supports = mlp(supports)

        assert (dim == dim_)

        logits = torch.zeros(self.num_ensemble, B, self.num_classes).to(z.device)
        for ens in range(self.num_ensemble):
            temp_centroids = (labels / (labels.sum(dim=0, keepdim=True) + 1e-12)).T @ mlp_supports[
                                                                                    ens * N: (ens + 1) * N]

            # normalize
            temp_z = torch.nn.functional.normalize(mlp_z[ens * B: (ens + 1) * B], dim=1)
            temp_centroids = torch.nn.functional.normalize(temp_centroids, dim=1)

            logits[ens] = self.tau * temp_z @ temp_centroids.T  # [B,C]

        return logits

    def select_supports(self):
        '''
        we filter support examples with high prediction entropy
        :return: filtered support examples.
        '''
        ent_s = self.ent
        y_hat = self.labels.argmax(dim=1).long()
        filter_K = self.filter_K
        if filter_K == -1:
            indices = torch.LongTensor(list(range(len(ent_s))))
        else:
            indices = []
            indices1 = torch.LongTensor(list(range(len(ent_s))))
            for i in range(self.num_classes):
                _, indices2 = torch.sort(ent_s[y_hat == i])
                indices2 = indices2.cpu()
                y_hat = y_hat.cpu()
                indices.append(indices1[y_hat == i][indices2][:filter_K])
            indices = torch.cat(indices)

        self.supports = self.supports[indices]
        self.labels = self.labels[indices]
        self.ent = self.ent[indices]

        return self.supports, self.labels

    @torch.enable_grad()
    def forward_and_adapt(self, z, supports, labels):
        # targets : pseudo labels, outputs: for prediction
        with torch.no_grad():
            targets, outputs = self.target_generation(z, supports, labels)

        self.optimizer.zero_grad()

        loss = None
        logits = self.compute_logits(z, supports, labels, self.mlps)

        for ens in range(self.num_ensemble):
            if loss is None:
                loss = F.kl_div(logits[ens].log_softmax(-1), targets[ens])
            else:
                loss += F.kl_div(logits[ens].log_softmax(-1), targets[ens])

        loss.backward()
        self.optimizer.step()

        return outputs  # outputs

    def reset(self):
        self.supports = self.warmup_supports.data
        self.labels = self.warmup_labels.data
        self.ent = self.warmup_ent.data

        self.mlps.reset()
        self.optimizer = torch.optim.Adam(self.mlps.parameters(), lr=self.lr)

        torch.cuda.empty_cache()

    # from https://jejjohnson.github.io/research_journal/snippets/numpy/euclidean/
    def euclidean_distance_einsum(self, X, Y):
        # X, Y [n, dim], [m, dim] -> [n,m]
        XX = torch.einsum('nd, nd->n', X, X)[:, None]  # [n, 1]
        YY = torch.einsum('md, md->m', Y, Y)  # [m]
        XY = 2 * torch.matmul(X, Y.T)  # [n,m]
        return XX + YY - XY

    def cosine_distance_einsum(self, X, Y):
        # X, Y [n, dim], [m, dim] -> [n,m]
        X = F.normalize(X, dim=1)
        Y = F.normalize(Y, dim=1)
        XX = torch.einsum('nd, nd->n', X, X)[:, None]  # [n, 1]
        YY = torch.einsum('md, md->m', Y, Y)  # [m]
        XY = 2 * torch.matmul(X, Y.T)  # [n,m]
        return XX + YY - XY

    def target_generation(self, z, supports, labels):
        # retrieve k nearest neighbors. from "https://github.com/csyanbin/TPN/blob/master/train.py"
        dist = self.cosine_distance_einsum(z, supports)
        W = torch.exp(-dist)  # [B, N]

        temp_k = self.filter_K if self.filter_K != -1 else supports.size(0) // self.num_classes
        k = min(self.k, temp_k)

        values, indices = torch.topk(W, k, sorted=False)  # [B, k]
        topk_indices = torch.zeros_like(W).scatter_(1, indices, 1)  # [B, N] 1 for topk, 0 for else
        temp_labels = self.compute_logits(supports, supports, labels, self.mlps)  # [ens, N, C]
        temp_labels_targets = F.one_hot(temp_labels.argmax(-1), num_classes=self.num_classes).float()  # [ens, N, C]
        temp_labels_outputs = torch.softmax(temp_labels, -1)  # [ens, N, C]

        topk_indices = topk_indices.unsqueeze(0).repeat(self.num_ensemble, 1, 1)  # [ens, B, N]

        # targets for pseudo labels. we use one-hot class distribution
        targets = torch.bmm(topk_indices, temp_labels_targets)  # [ens, B, C]
        targets = targets / (targets.sum(2, keepdim=True) + 1e-12)  # [ens, B, C]
        #targets = targets.mean(0)  # [B,C]

        # outputs for prediction
        outputs = torch.bmm(topk_indices, temp_labels_outputs)  # [ens, B, C]
        outputs = outputs / (outputs.sum(2, keepdim=True) + 1e-12)  # [ens, B, C]
        outputs = outputs.mean(0)  # [B,C]

        return targets, outputs


class TAST_BN(nn.Module):
    def __init__(self, algorithm, lr=5e-3):
        super().__init__()
        self.featurizer = algorithm.featurizer
        self.classifier = algorithm.classifier

        # store supports examples and corresponding labels. we store test example itself instead of the feature representations.
        self.supports = None
        self.labels = None
        self.ent = None

        self.num_classes = self.classifier.fc.weight.data.size()[0]

        # hparams
        # adapt_hparams_dict = {
        #     'filter_K': [1, 5, 20],
        #     'gamma': [1, 3],
        #     'lr': [1e-3],
        #     'tau': [10],
        #     'k': [1, 2, 4, 8],
        # }
        self.filter_K = 2  # filter_K > 3 will out of memory

        # we restrict the size of support set
        if self.filter_K * self.num_classes >150 :
            self.filter_K = int(150 / self.num_classes)

        self.steps = 1
        self.lr = lr
        self.tau = 10
        # self.num_classes = num_classes
        self.k = 4

        self.model, self.optimizer = self.configure_model_optimizer(algorithm)

        self.model_state, self.optimizer_state = \
            self.copy_model_and_optimizer(self.model, self.optimizer)


    def copy_model_and_optimizer(self, model, optimizer):
        """Copy the model and optimizer states for resetting after adaptation."""
        model_state = copy.deepcopy(model.state_dict())
        optimizer_state = copy.deepcopy(optimizer.state_dict())
        return model_state, optimizer_state

    def load_model_and_optimizer(self, model, optimizer, model_state, optimizer_state):
        """Restore the model and optimizer states from copies."""
        model.load_state_dict(model_state, strict=True)
        optimizer.load_state_dict(optimizer_state)

    def collect_params(self, model):
        params = []
        names = []
        for nm, m in model.named_modules():
            if isinstance(m, nn.BatchNorm2d):
                for np, p in m.named_parameters():
                    if np in ['weight', 'bias']:
                        params.append(p)
                        names.append(f"{nm}.{np}")
        return params, names

    def configure_model(self, model):
        model.train()
        model.requires_grad_(False)
        for m in model.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.requires_grad_(True)
                m.track_running_stats = False
                m.running_mean = None
                m.running_var = None
        return model

    def configure_model_optimizer(self, algorithm):
        adapted_algorithm = copy.deepcopy(algorithm)
        adapted_algorithm.featurizer = self.configure_model(adapted_algorithm.featurizer)
        params, param_names = self.collect_params(adapted_algorithm.featurizer)
        optimizer = torch.optim.Adam(params, lr=self.lr)

        return adapted_algorithm, optimizer

    def forward(self, x, adapt=True):
        if adapt:
            p_supports = self.model.predict(x)
            # p_supports = self.model(x)
            yhat = torch.nn.functional.one_hot(p_supports.argmax(1), num_classes=self.num_classes).float()
            ent = softmax_entropy(p_supports)

            # prediction

            if self.supports is None:
                self.supports = x.detach()
                self.labels = yhat.detach()
                self.ent = ent.detach()
            else:
                self.supports = self.supports.to(x.device)
                self.labels = self.labels.to(x.device)
                self.ent = self.ent.to(x.device)
                self.supports = torch.cat([self.supports, x.detach()])
                self.labels = torch.cat([self.labels, yhat.detach()])
                self.ent = torch.cat([self.ent, ent.detach()])

        supports, labels = self.select_supports()

        for _ in range(self.steps):
            p = self.forward_and_adapt(x, supports, labels)

        return p

    def compute_logits(self, z, supports, labels):
        B, dim = z.size()
        N, dim_ = supports.size()

        temp_centroids = (labels / (labels.sum(dim=0, keepdim=True) + 1e-12)).T @ supports

        # normalize
        temp_z = torch.nn.functional.normalize(z, dim=1)
        temp_centroids = torch.nn.functional.normalize(temp_centroids, dim=1)

        logits = self.tau * temp_z @ temp_centroids.T  # [B,C]

        return logits

    def select_supports(self):
        ent_s = self.ent
        y_hat = self.labels.argmax(dim=1).long()
        filter_K = self.filter_K
        if filter_K == -1:
            indices = torch.LongTensor(list(range(len(ent_s))))
        else:
            indices = []
            indices1 = torch.LongTensor(list(range(len(ent_s))))
            for i in range(self.num_classes):
                _, indices2 = torch.sort(ent_s[y_hat == i])
                indices2 = indices2.cpu()
                y_hat = y_hat.cpu()
                indices.append(indices1[y_hat == i][indices2][:filter_K])
            indices = torch.cat(indices)

        self.supports = self.supports[indices]
        self.labels = self.labels[indices]
        self.ent = self.ent[indices]

        return self.supports, self.labels

    @torch.enable_grad()
    def forward_and_adapt(self, x, supports, labels):
        feats = torch.cat((x, supports), 0)
        feats = self.model.featurizer(feats)
        z, supports = feats[:x.size(0)], feats[x.size(0):]
        del feats

        with torch.no_grad():
            targets, outputs = self.target_generation(z, supports, labels)

        self.optimizer.zero_grad()

        # PL with Ensemble
        logits = self.compute_logits(z, supports, labels)
        loss = F.kl_div(logits.log_softmax(-1), targets)

        loss.backward()
        self.optimizer.step()

        return outputs

    def reset(self):
        self.supports = None
        self.labels = None
        self.ent = None
        self.load_model_and_optimizer(self.model, self.optimizer,
                                      self.model_state, self.optimizer_state)

    def euclidean_distance_einsum(self, X, Y):
        # X, Y [n, dim], [m, dim] -> [n,m]
        XX = torch.einsum('nd, nd->n', X, X)[:, None]  # [n, 1]
        YY = torch.einsum('md, md->m', Y, Y)  # [m]
        XY = 2 * torch.matmul(X, Y.T)  # [n,m]
        return XX + YY - XY

    def cosine_distance_einsum(self, X, Y):
        # X, Y [n, dim], [m, dim] -> [n,m]
        X = F.normalize(X, dim=1)
        Y = F.normalize(Y, dim=1)
        XX = torch.einsum('nd, nd->n', X, X)[:, None]  # [n, 1]
        YY = torch.einsum('md, md->m', Y, Y)  # [m]
        XY = 2 * torch.matmul(X, Y.T)  # [n,m]
        return XX + YY - XY

    def target_generation(self, z, supports, labels):
        dist = self.cosine_distance_einsum(z, supports)
        W = torch.exp(-dist)  # [B, N]

        temp_k = self.filter_K if self.filter_K != -1 else supports.size(0) // self.num_classes
        k = min(self.k, temp_k)

        values, indices = torch.topk(W, k, sorted=False)  # [B, k]
        topk_indices = torch.zeros_like(W).scatter_(1, indices, 1)  # [B, N] 1 for topk, 0 for else
        temp_labels = self.compute_logits(supports, supports, labels)  # [N, C]
        temp_labels_targets = F.one_hot(temp_labels.argmax(-1), num_classes=self.num_classes).float()  # [N, C]
        temp_labels_outputs = torch.softmax(temp_labels, -1)  # [N, C]

        targets = topk_indices @ temp_labels_targets
        outputs = topk_indices @ temp_labels_outputs

        targets = targets / (targets.sum(-1, keepdim=True) + 1e-12)
        outputs = outputs / (outputs.sum(-1, keepdim=True) + 1e-12)

        return targets, outputs
    

class DeYO(nn.Module):
    def __init__(self, model, optimizer, steps=1, episodic=False):
        super().__init__()
        self.model = model
        self.optimizer = optimizer
        self.steps = steps
        assert steps > 0, "tent requires >= 1 step(s) to forward and update"
        self.episodic = episodic
        
        self.ent_threshold = 0.5  # Entropy threshold for sample selection $\tau_\mathrm{Ent}$ in Eqn. (8)
        self.plpd_threshold = 0.3  # PLPD threshold for sample selection $\tau_\mathrm{PLPD}$ in Eqn. (8)
        self.ent0 = 0.4  # Entropy margin for sample weighting $\mathrm{Ent}_0$ in Eqn. (10)
        self.patch_len = 4  # The number of patches per row/column

    def forward(self, x):
        if self.episodic:
            self.reset()

        for _ in range(self.steps):
            outputs = self.forward_and_adapt(x, self.model, self.optimizer)

        return outputs

    def reset(self):
        if self.model_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer state")
        load_model_and_optimizer(self.model, self.optimizer,
                                 self.model_state, self.optimizer_state)
    
    @torch.enable_grad()
    def forward_and_adapt(self,x, model, optimizer):
        """Forward and adapt model input data.
        Measure entropy of the model prediction, take gradients, and update params.
        """
        outputs = model.predict(x)
        
        optimizer.zero_grad()
        entropys = softmax_entropy(outputs)

        filter_ids_1 = torch.where((entropys < self.ent_threshold))

        entropys = entropys[filter_ids_1]
        backward = len(entropys)
        if backward==0:
            return outputs

        x_prime = x[filter_ids_1]
        x_prime = x_prime.detach()

        # aug_type=='patch':
        resize_t = torchvision.transforms.Resize(((x.shape[-1]//self.patch_len)*self.patch_len,(x.shape[-1]//self.patch_len)*self.patch_len))
        resize_o = torchvision.transforms.Resize((x.shape[-1],x.shape[-1]))
        x_prime = resize_t(x_prime)
        x_prime = rearrange(x_prime, 'b c (ps1 h) (ps2 w) -> b (ps1 ps2) c h w', ps1=self.patch_len, ps2=self.patch_len)
        perm_idx = torch.argsort(torch.rand(x_prime.shape[0],x_prime.shape[1]), dim=-1)
        x_prime = x_prime[torch.arange(x_prime.shape[0]).unsqueeze(-1),perm_idx]
        x_prime = rearrange(x_prime, 'b (ps1 ps2) c h w -> b c (ps1 h) (ps2 w)', ps1=self.patch_len, ps2=self.patch_len)
        x_prime = resize_o(x_prime)
        
        with torch.no_grad():
            outputs_prime = model.predict(x_prime)
        
        prob_outputs = outputs[filter_ids_1].softmax(1)
        prob_outputs_prime = outputs_prime.softmax(1)

        cls1 = prob_outputs.argmax(dim=1)

        plpd = torch.gather(prob_outputs, dim=1, index=cls1.reshape(-1,1)) - torch.gather(prob_outputs_prime, dim=1, index=cls1.reshape(-1,1))
        plpd = plpd.reshape(-1)
        
        
        filter_ids_2 = torch.where(plpd > self.plpd_threshold)

        entropys = entropys[filter_ids_2]
        final_backward = len(entropys)
            
        if final_backward==0:
            del x_prime
            del plpd
            return outputs
            
        plpd = plpd[filter_ids_2]

        coeff = ((1 / (torch.exp(((entropys.clone().detach()) - self.ent0)))) +
                (1 / (torch.exp(-1. * plpd.clone().detach())))
                )            
        entropys = entropys.mul(coeff)
        loss = entropys.mean(0)

        if final_backward != 0:
            loss.backward()
            optimizer.step()
        optimizer.zero_grad()

        del x_prime
        del plpd
        
        return outputs